package top.continew.ai.service.impl;

import cn.hutool.core.collection.CollectionUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import top.continew.ai.service.VectorStoreService;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @Author： Lonni
 * @Date： 2025/8/21  15:15
 */
@Service
@Slf4j
public class VectorStoreServiceImpl implements VectorStoreService {

    /**
     * Spring AI提供的向量存储接口
     */
    @Autowired
    private VectorStore vectorStore;

    /**
     * 增加文档
     *
     * @param documents
     * @return
     */
    @Override
    public Boolean addDocuments(List<Document> documents) {
        vectorStore.add(documents);
        return true;
    }

    /**
     * 相似性检索
     *
     * @param query
     * @return
     */
    @Override
    public List<Document> retrieveDocuments(String query, int topK, double scoreThreshold) {
        SearchRequest build = SearchRequest.builder()
                .query(query)
                .topK(topK)
                .similarityThreshold(scoreThreshold).build();
        return vectorStore.similaritySearch(build);
    }

    /**
     * 根据元数据删除文档
     *
     * @param key
     * @param value
     * @return
     */
    @Override
    public Boolean deleteDocumentsByMetadata(String key, String value) {
        //先查询源数据
        List<Document> documents = vectorStore.similaritySearch(SearchRequest.builder().query("").filterExpression(key + "==" + value).build());
        if (CollectionUtil.isNotEmpty(documents)) {
            vectorStore.delete(documents.stream().map(Document::getId).collect(Collectors.toList()));
        }
        return true;
    }
}
